package com.example.shardingjdbc.sharding.algorithm;

import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.util.*;

/**
 * @Description: 复合分片算法
 * @Author: Jerry
 * @Date: 2021-07-21 22:18
 */
public class MyComplexShardingAlgorithm implements ComplexKeysShardingAlgorithm<Long> {
    private final static String ORDER_ID = "order_id";
    private final static String USER_ID = "user_id";

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames,
                                         ComplexKeysShardingValue<Long> shardingValue) {

        Collection<Long> orderIds = getShardingValue(shardingValue, ORDER_ID);
        Collection<Long> userIds = getShardingValue(shardingValue, USER_ID);

        Set<String> tables = new HashSet<>();
        for (Long userId : userIds) {
            for (Long orderId : orderIds) {
                int index = getIndex(userId, orderId, availableTargetNames);
                for (String tableName: availableTargetNames) {
                    if (tableName.endsWith(String.valueOf(index))) {
                        tables.add(tableName);
                    }
                }
            }
        }

        return tables;
    }

    private Collection<Long> getShardingValue(ComplexKeysShardingValue shardingValue, String key) {
        Collection<Long> valueList = new ArrayList<>();
        Map<String, Collection<Long>> columnNameAndShardingValuesMap =
                shardingValue.getColumnNameAndShardingValuesMap();
        if (columnNameAndShardingValuesMap.containsKey(key)) {
            valueList.addAll(columnNameAndShardingValuesMap.get(key));
        }

        return valueList;
    }

    private int getIndex(Long userId, Long orderId, Collection<String> availableTargetNames) {
        String userIdStr = String.valueOf(userId);
        String orderIdStr = String.valueOf(orderId);
        int sum = Integer.parseInt(userIdStr.substring(userIdStr.length() - 1)) + Integer.parseInt(orderIdStr.substring(orderIdStr.length() - 1));
        return sum % availableTargetNames.size();
    }
}
